package com.mxx.common.data.jpa;

import com.mxx.common.data.utils.EntityReflectInvoke;
import com.mxx.common.data.utils.SortPaged;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.data.jpa.repository.support.JpaEntityInformation;
import org.springframework.data.repository.support.PageableExecutionUtils;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

import javax.persistence.EntityManager;
import javax.persistence.TypedQuery;
import javax.persistence.criteria.*;
import java.util.*;

public class JpaOperationUtil {

    public static <T> Object getResult(EntityManager entityManager, Class domainClass, TypedQuery<T> typedQuery, Specification<T> spec, Pageable pageable) {
        if (pageable instanceof SortPaged) {
            return typedQuery.getResultList();
        } else {
            return (pageable.isUnpaged() ? new PageImpl<T>(typedQuery.getResultList()) : readPage(entityManager, typedQuery, domainClass, pageable, spec));
        }
    }

    public static Object toDomainObject(Class domainClass, Class voClass, String[] columns, Object[] jpaData) {
        Object entity = EntityReflectInvoke.newInterface(voClass != null ? voClass : domainClass);
        EntityReflectInvoke.set(entity, columns, jpaData);
        return entity;
    }


    public static <S> Page<S> readPage(EntityManager entityManager, TypedQuery<S> query, final Class<S> domainClass, Pageable pageable,
                                       @Nullable Specification<S> spec) {

        if (pageable.isPaged()) {
            query.setFirstResult((int) pageable.getOffset());
            query.setMaxResults(pageable.getPageSize());
        }

        return PageableExecutionUtils.getPage(query.getResultList(), pageable,
                () -> executeCountQuery(getCountQuery(entityManager, spec, domainClass)));
    }

    private static long executeCountQuery(TypedQuery<Long> query) {

        Assert.notNull(query, "TypedQuery must not be null!");

        List<Long> totals = query.getResultList();
        long total = 0L;

        for (Long element : totals) {
            total += element == null ? 0 : element;
        }

        return total;
    }

    public static <S> TypedQuery<Long> getCountQuery(EntityManager entityManager, @Nullable Specification<S> spec, Class<S> domainClass) {

        CriteriaBuilder builder = entityManager.getCriteriaBuilder();
        CriteriaQuery<Long> query = builder.createQuery(Long.class);

        Root<S> root = applySpecificationToCriteria(entityManager, spec, domainClass, query);

        if (query.isDistinct()) {
            query.select(builder.countDistinct(root));
        } else {
            query.select(builder.count(root));
        }

        // Remove all Orders the Specifications might have applied
        query.orderBy(Collections.<Order>emptyList());

        return entityManager.createQuery(query);
    }

    public static <S, U> Root<U> applySpecificationToCriteria(EntityManager entityManager, @Nullable Specification<U> spec, Class<U> domainClass,
                                                              CriteriaQuery<S> query) {

        Assert.notNull(domainClass, "Domain class must not be null!");
        Assert.notNull(query, "CriteriaQuery must not be null!");

        Root<U> root = query.from(domainClass);

        if (spec == null) {
            return root;
        }

        CriteriaBuilder builder = entityManager.getCriteriaBuilder();
        Predicate predicate = spec.toPredicate(root, query, builder);

        if (predicate != null) {
            query.where(predicate);
        }

        return root;
    }

    public static <T> void selector(String[] columns, Root<T> root, CriteriaQuery<T> criteriaQuery) {
        List<Selection<?>> selections = new ArrayList<>(columns.length);
        for (int index = 0; index < columns.length; index++) {
            String column = columns[index];
            selections.add(root.get(column));
        }
        criteriaQuery.multiselect(selections);
    }
}
